package com.pai4j.zwiki.websocket.handler;

import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * ZWiki通知WebSocket处理器
 * 
 * 核心功能：
 * 1. 管理用户的WebSocket连接会话
 * 2. 处理客户端发送的消息（如心跳检测）
 * 3. 向客户端推送实时通知
 * 4. 处理连接建立、断开、异常
 *
 * 会话管理：
 * - 使用ConcurrentHashMap存储所有用户的WebSocket会话
 * - KEY: userId（用户ID）
 * - VALUE: WebSocketSession（WebSocket会话对象）
 *
 * @author pai4j
 * @since 2025-10-14
 */
@Slf4j
public class ZWikiNotificationWebSocketHandler extends TextWebSocketHandler {

    /**
     * 存储所有用户的WebSocket会话
     * KEY: userId（用户ID）
     * VALUE: WebSocketSession（WebSocket会话对象）
     * 
     * 线程安全：使用ConcurrentHashMap保证并发安全
     */
    public static final Map<String, WebSocketSession> USER_SESSIONS = new ConcurrentHashMap<>();

    /**
     * 连接建立成功后的回调
     * 
     * @param session WebSocket会话
     * @throws Exception 异常
     */
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        // 从会话属性中获取userId（由拦截器设置）
        Map<String, Object> attributes = session.getAttributes();
        String userId = (String) attributes.get("userId");

        if (StringUtils.isBlank(userId)) {
            log.warn("WebSocket连接缺少userId参数，关闭连接. sessionId: {}", session.getId());
            session.close(CloseStatus.BAD_DATA);
            return;
        }

        // 如果用户已有连接，关闭旧连接，防止重复连接堆积
        WebSocketSession oldSession = USER_SESSIONS.get(userId);
        if (oldSession != null && oldSession.isOpen() && !StringUtils.equals(oldSession.getId(), session.getId())) {
            try {
                log.warn("检测到重复的WebSocket连接，关闭旧连接. userId: {}, oldSessionId: {}, newSessionId: {}", 
                        userId, oldSession.getId(), session.getId());
                oldSession.close(CloseStatus.NORMAL);
            } catch (Exception e) {
                log.warn("关闭旧WebSocket会话失败. userId: {}, oldSessionId: {}", 
                        userId, oldSession.getId(), e);
            }
        }

        // 保存新连接
        USER_SESSIONS.put(userId, session);

        log.info("✅ ZWiki WebSocket连接建立成功. userId: {}, sessionId: {}, 当前在线用户数: {}", 
                userId, session.getId(), USER_SESSIONS.size());

        // 发送欢迎消息
        sendWelcomeMessage(session, userId);
    }

    /**
     * 接收客户端发送的文本消息
     * 
     * @param session WebSocket会话
     * @param message 文本消息
     * @throws Exception 异常
     */
    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        String payload = message.getPayload();
        
        log.debug("收到WebSocket消息: {}", payload);

        // 处理心跳检测
        if ("ping".equals(payload)) {
            session.sendMessage(new TextMessage("pong"));
            return;
        }

        // 处理订阅请求
        if (payload.startsWith("subscribe:")) {
            String taskId = payload.substring("subscribe:".length());
            // 保存订阅信息到会话属性
            session.getAttributes().put("subscribedTaskId", taskId);
            
            String response = String.format("{\"type\":\"subscribed\",\"taskId\":\"%s\",\"message\":\"订阅成功\"}", taskId);
            session.sendMessage(new TextMessage(response));
            
            log.info("用户订阅任务: userId={}, taskId={}", 
                    session.getAttributes().get("userId"), taskId);
            return;
        }

        // 处理取消订阅
        if (payload.startsWith("unsubscribe:")) {
            session.getAttributes().remove("subscribedTaskId");
            session.sendMessage(new TextMessage("{\"type\":\"unsubscribed\",\"message\":\"取消订阅成功\"}"));
            return;
        }

        // 其他消息类型暂不处理
        log.debug("未识别的消息类型: {}", payload);
    }

    /**
     * 连接关闭后的回调
     * 
     * @param session WebSocket会话
     * @param status 关闭状态
     * @throws Exception 异常
     */
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        Map<String, Object> attributes = session.getAttributes();
        String userId = (String) attributes.get("userId");

        if (StringUtils.isNotBlank(userId)) {
            // 移除会话（只移除ID匹配的会话，防止误删新连接）
            WebSocketSession existing = USER_SESSIONS.get(userId);
            if (existing != null && StringUtils.equals(existing.getId(), session.getId())) {
                USER_SESSIONS.remove(userId);
            }
        }

        log.info("❌ ZWiki WebSocket连接关闭. userId: {}, sessionId: {}, status: {}, 当前在线用户数: {}", 
                userId, session.getId(), status, USER_SESSIONS.size());
    }

    /**
     * 传输错误处理
     * 
     * @param session WebSocket会话
     * @param exception 异常
     * @throws Exception 异常
     */
    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        Map<String, Object> attributes = session.getAttributes();
        String userId = (String) attributes.get("userId");

        log.error("⚠️ ZWiki WebSocket传输错误. userId: {}, sessionId: {}", 
                userId, session.getId(), exception);

        // 尝试关闭会话
        try {
            if (session.isOpen()) {
                session.close(CloseStatus.SERVER_ERROR);
            }
        } catch (Exception e) {
            log.debug("关闭异常会话失败. sessionId: {}", session.getId(), e);
        }
    }

    /**
     * 发送欢迎消息
     */
    private void sendWelcomeMessage(WebSocketSession session, String userId) {
        try {
            String welcomeMsg = String.format(
                    "{\"type\":\"welcome\",\"userId\":\"%s\",\"message\":\"欢迎连接ZWiki实时通知服务！\"," +
                    "\"features\":[\"任务状态更新\",\"队列位置通知\",\"进度实时推送\",\"完成/失败通知\"]}",
                    userId
            );
            session.sendMessage(new TextMessage(welcomeMsg));
        } catch (Exception e) {
            log.warn("发送欢迎消息失败. userId: {}", userId, e);
        }
    }

    // ========== 公共方法：供其他服务调用 ==========

    /**
     * 获取用户的WebSocket会话
     *
     * @param userId 用户ID
     * @return WebSocket会话，如果用户未连接则返回null
     */
    public static WebSocketSession getSession(String userId) {
        if (StringUtils.isBlank(userId)) {
            return null;
        }
        return USER_SESSIONS.get(userId);
    }

    /**
     * 向指定用户发送消息
     *
     * @param userId 用户ID
     * @param message 消息内容（JSON格式）
     * @return 是否发送成功
     */
    public static boolean sendToUser(String userId, String message) {
        if (StringUtils.isBlank(userId) || StringUtils.isBlank(message)) {
            return false;
        }

        WebSocketSession session = getSession(userId);
        if (session == null || !session.isOpen()) {
            log.debug("用户WebSocket未连接或已关闭. userId: {}", userId);
            return false;
        }

        try {
            session.sendMessage(new TextMessage(message));
            log.debug("✉️ 成功推送消息到用户. userId: {}, messageLength: {}", userId, message.length());
            return true;
        } catch (IOException e) {
            log.error("推送消息失败. userId: {}", userId, e);
            return false;
        }
    }

    /**
     * 向所有在线用户广播消息
     *
     * @param message 消息内容
     * @return 成功推送的用户数
     */
    public static int broadcastToAll(String message) {
        int successCount = 0;

        for (Map.Entry<String, WebSocketSession> entry : USER_SESSIONS.entrySet()) {
            String userId = entry.getKey();
            WebSocketSession session = entry.getValue();

            if (session != null && session.isOpen()) {
                try {
                    session.sendMessage(new TextMessage(message));
                    successCount++;
                } catch (IOException e) {
                    log.warn("广播消息失败. userId: {}", userId, e);
                }
            }
        }

        log.info("📢 广播消息完成. 在线用户: {}, 成功推送: {}", USER_SESSIONS.size(), successCount);
        return successCount;
    }

    /**
     * 获取在线用户数
     *
     * @return 在线用户数
     */
    public static int getOnlineUserCount() {
        return USER_SESSIONS.size();
    }

    /**
     * 检查用户是否在线
     *
     * @param userId 用户ID
     * @return 是否在线
     */
    public static boolean isUserOnline(String userId) {
        if (StringUtils.isBlank(userId)) {
            return false;
        }
        WebSocketSession session = USER_SESSIONS.get(userId);
        return session != null && session.isOpen();
    }

    /**
     * 关闭指定用户的连接
     *
     * @param userId 用户ID
     * @param reason 关闭原因
     */
    public static void closeUserConnection(String userId, String reason) {
        WebSocketSession session = USER_SESSIONS.get(userId);
        if (session != null && session.isOpen()) {
            try {
                session.close(new CloseStatus(1000, reason));
                USER_SESSIONS.remove(userId);
                log.info("关闭用户WebSocket连接. userId: {}, reason: {}", userId, reason);
            } catch (Exception e) {
                log.error("关闭WebSocket连接失败. userId: {}", userId, e);
            }
        }
    }
}

